{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0c3d63fc",
   "metadata": {},
   "source": [
    "# QM/MM simulation: VQE within a Polarizable Embedded Framework.\n",
    "\n",
    "In this tutorial, we will discuss a hybrid quantum-classical algorithm called PE-VQE-SCF introduced in this [paper](https://arxiv.org/pdf/2312.01926). The PE-VQE-scf approch combines:\n",
    "\n",
    "- Variational Quantum Eigensolver (VQE) – a quantum algorithm for finding ground-state energies of molecules.\n",
    "- Self-Consistent Field (SCF) – a classical iterative method used in quantum chemistry.\n",
    "- Polarizable Embedding (PE) – a technique to model the effect of an environment (like a solvent or protein) on a quantum system.\n",
    "\n",
    "The goal is to simulate chemical systems on quantum computers (or simulator) by embedding them in a polarizable environment. In this tutorial, we will use CUDA-Q to implement the QM/MM framework and the CUDA-Q GPU accelerated simulator for running the simulation.\n",
    "\n",
    "## Key concepts:\n",
    "\n",
    "1- Variational Quantum Eigensolver (VQE)\n",
    "- A quantum algorithm that estimates the ground state energy of a molecule.\n",
    "- It uses a parameterized quantum circuit (ansatz) and a classical optimizer to minimize the energy expectation value.\n",
    "- In this tutorial, we employ VQE with UCCSD ansatz for the quantum part. However, user should be able to replace the VQE and the ansatz with other quantum algorithm.\n",
    "\n",
    "2- Self-Consistent Field (SCF)\n",
    "- A method where the solution (e.g., molecular orbitals) is updated iteratively until convergence.\n",
    "- In this context, SCF is used to update the embedding potential based on the quantum system's density.\n",
    "\n",
    "3- Polarizable Embedding (PE)\n",
    "- Models the environment as a set of polarizable sites that respond to the quantum system's electric field.\n",
    "- The environment affects the quantum system, and vice versa, requiring mutual polarization.\n",
    "\n",
    "## PE-VQE-SCF Algorithm Steps\n",
    "<div>\n",
    "<img src=\"images/qm-mm-pe.png\" width=\"600\">\n",
    "</div>\n",
    "\n",
    "### Step 1: Initialize (Classical pre-processing)\n",
    "- Define the quantum region (e.g., a molecule) and the classical environment (e.g., solvent).\n",
    "- Set up the PE parameters (multipoles, polarizabilities).\n",
    "\n",
    "### Step 2: Build the Hamiltonian\n",
    "- Construct the PE-augmented Hamiltonian that includes the molecular Hamiltonian and the interaction with the polarizable environment.\n",
    "\n",
    "### Step 3: Run VQE\n",
    "- Use a quantum computer (or simulator) to estimate the ground state energy of the system using VQE.\n",
    "- The quantum circuit prepares a trial wavefunction, and the classical optimizer adjusts parameters to minimize energy.\n",
    "\n",
    "### Step 4: Update Environment\n",
    "- Compute the electronic density from the VQE result.\n",
    "- Update the induced dipoles in the environment based on this density.\n",
    "\n",
    "### Step 5: Self-Consistency Loop\n",
    "- Repeat Steps 2–4 until the energy and density converge (i.e., SCF convergence)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93fe0a91",
   "metadata": {},
   "source": [
    "### Requirments:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae619fe8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# You might need to install (apt-get update & apt-get install build-essential -y \n",
    "#                         apt-get install python3-dev) before installing cppe.\n",
    "%pip install cppe\n",
    "%pip install pyscf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "af9dd1da",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyscf import gto, scf, solvent\n",
    "import numpy as np\n",
    "from functools import reduce\n",
    "\n",
    "import cudaq\n",
    "\n",
    "cudaq.set_target(\"nvidia\", option=\"fp64\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42de878a",
   "metadata": {},
   "source": [
    "### Example 1: LiH with 2 water molecules."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a73e83dc",
   "metadata": {},
   "source": [
    "#### Initialize and build the hamiltonian."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "5ee6dbc8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "overwrite output file: Li 0-pyscf.log\n",
      "[Pyscf] Total number of electrons =  4\n",
      "[Pyscf] Total number of orbitals =  6\n",
      "[Pyscf] Total HF energy with solvent: -7.862708330327407\n",
      "[Pyscf] Polarizable embedding energy from HF:  -1.1127389358871828e-05\n",
      "[Pyscf] Total CCSD energy with solvent:  -7.882680504632604\n",
      "[Pyscf] Polarizable embedding energy from CCSD:  -1.0898134916648752e-05\n",
      "        Running solver for induced moments.\n",
      "0        --- Norm: 0.000095431562\n",
      "1        --- Norm: 0.000000927586\n",
      "        --- Turning on DIIS. ---\n",
      "2        --- Norm: 0.000000010648\n",
      "3        --- Norm: 0.000000000133\n"
     ]
    }
   ],
   "source": [
    "from qchem.classical_pyscf import get_mol_pe_hamiltonian\n",
    "\n",
    "geometry = 'Li 0.3925 0.0 0.0; H -1.1774 0.0 0.0'\n",
    "water_pot= \"qchem/4NP_in_water.pot\"\n",
    "\n",
    "molecular_data=get_mol_pe_hamiltonian(xyz=geometry, potfile=water_pot, spin=0, charge=0, basis='sto3g', ccsd=True, verbose=True)\n",
    "\n",
    "obi_mol = molecular_data[0]\n",
    "tbi_mol = molecular_data[1]\n",
    "e_nn = molecular_data[2]\n",
    "obi_pe = molecular_data[3]\n",
    "nelectrons = molecular_data[4]\n",
    "norbitals = molecular_data[5]\n",
    "\n",
    "qubits_num = 2 * norbitals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "fd32d51b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from qchem.hamiltonian import jordan_wigner_fermion, jordan_wigner_pe\n",
    "\n",
    "spin_mol_ham=jordan_wigner_fermion(obi_mol, tbi_mol, e_nn, tolerance = 1e-12)\n",
    "\n",
    "spin_pe_ham=jordan_wigner_pe(obi_pe, tolerance = 1e-12)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "52780519",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "converged SCF energy = -7.86270833032741\n"
     ]
    }
   ],
   "source": [
    "mol = gto.M(\n",
    "    atom = geometry,\n",
    "    spin = 0,\n",
    "    charge = 0,\n",
    "    basis = 'sto3g'\n",
    ")\n",
    "\n",
    "mf = scf.RHF(mol)\n",
    "mf_pe = solvent.PE(mf, water_pot).run()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8932c03",
   "metadata": {},
   "source": [
    "### VQE, update environment, and scf loop.\n",
    "\n",
    "Note: for more accurate result, user needs to change `conv_tol` to much smaller value, e.g., 1e-7. In addition, the tolerance for the classical optimizer `vqe_tol` should be modified to smaller value, e.g., 1e-7."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d4f398a6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of parameters: 16 singles, 76 doubles, 92 total\n",
      "overwrite output file: Li 0-pyscf.log\n",
      "Optimizer exited successfully:  True\n",
      "VQE Energy: -7.867161909761761\n",
      "VQE converged: True\n",
      "        Running solver for induced moments.\n",
      "0        --- Norm: 0.000095457621\n",
      "1        --- Norm: 0.000000860362\n",
      "        --- Turning on DIIS. ---\n",
      "2        --- Norm: 0.000000004312\n",
      "Energy from observe call=     0.000762169 & Energy from np einsum=     0.000762197\n",
      "Cycle 1  E_tot = -7.86793486027634  dE = 7.86793\n",
      "Polarizable embedding energy:  -1.0781673490428836e-05\n",
      "Expectation value of polarizable operator:  0.0007621688410912728\n",
      "\n",
      "\n",
      "Optimizer exited successfully:  True\n",
      "VQE Energy: -7.8692404296251395\n",
      "VQE converged: True\n",
      "        Running solver for induced moments.\n",
      "0        --- Norm: 0.000095564807\n",
      "1        --- Norm: 0.000000887381\n",
      "        --- Turning on DIIS. ---\n",
      "2        --- Norm: 0.000000007916\n",
      "Energy from observe call=     0.000762218 & Energy from np einsum=     0.000762229\n",
      "Cycle 2  E_tot = -7.87001339604889  dE = 0.00207854\n",
      "Polarizable embedding energy:  -1.0748496723151968e-05\n",
      "Expectation value of polarizable operator:  0.0007622179270272233\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "Final result: \n",
      "Polarizable embedding energy:  -1.0748496723151968e-05\n",
      "Expectation value of polarizable operator:  0.0007622179270272233\n",
      "PE-VQE-UCCSD energy (L-BFGS-B)=  -7.87001339604889\n",
      "Natural orbitals occupation number:  [1.99899690e+00 1.91601348e+00 6.71542420e-02 1.60555115e-02\n",
      " 1.15390205e-03 6.30790416e-04]\n"
     ]
    }
   ],
   "source": [
    "from qchem.uccsd import uccsd_parameter_size\n",
    "from qchem.uccsd_vqe import uccsd_circuit_vqe\n",
    "from qchem.particle_operator import one_particle_op\n",
    "from qchem.PEoperator import pe_operator\n",
    "from qchem.uccsd_init_param import get_parameters\n",
    "\n",
    "e_last = 0.0\n",
    "conv_tol = 1e-2\n",
    "dE = 1.0\n",
    "cycle = 1\n",
    "vqe_tol = 1e-4\n",
    "\n",
    "# Get the number of parameters\n",
    "singles, doubles, total = uccsd_parameter_size(nelectrons, qubits_num, spin_mult = 0)\n",
    "print(f\"Number of parameters: {singles} singles, {doubles} doubles, {total} total\")\n",
    "\n",
    "# Get the initial parameters for the UCCSD circuit\n",
    "theta = get_parameters(xyz=geometry, spin=0, charge=0, basis='sto3g', ccsd=True, without_solvent=False, potfile=water_pot)\n",
    "\n",
    "#theta = np.zeros(total, dtype=np.float64)\n",
    "\n",
    "\n",
    "while dE > conv_tol:\n",
    "#for i in range (2):\n",
    "    spin_ham = spin_mol_ham + spin_pe_ham\n",
    "    \n",
    "    vqe_result = uccsd_circuit_vqe(spin_mult = 0, only_singles = False, only_doubles = False, qubits_num = qubits_num, \n",
    "                                   electron_count = nelectrons, optimize = True, theta = theta, \n",
    "                                   hamiltonian = spin_ham, method = 'L-BFGS-B', vqe_tol = vqe_tol, verbose = False)\n",
    "    \n",
    "    print(f\"VQE Energy: {vqe_result[0]}\")\n",
    "    #print(f\"VQE parameters: {vqe_result[1]}\")\n",
    "    print(f\"VQE converged: {vqe_result[2]}\")\n",
    "    \n",
    "    \n",
    "    theta = vqe_result[1]\n",
    "    \n",
    "    if vqe_result[2]:\n",
    "\n",
    "        # Compute rdm in atomic orbital.\n",
    "        dm_a = np.zeros((qubits_num // 2, qubits_num // 2))\n",
    "        dm_b = np.zeros((qubits_num // 2, qubits_num // 2))\n",
    "            \n",
    "        n_spatial_orbitals = qubits_num // 2\n",
    "            \n",
    "        n_occupied = nelectrons // 2\n",
    "        n_virtual = n_spatial_orbitals - n_occupied\n",
    "\n",
    "        alpha_indices = [i * 2 for i in range(n_occupied)]\n",
    "        alpha_indices += [i * 2 + nelectrons for i in range(n_virtual)]\n",
    "\n",
    "        beta_indices = [i * 2 + 1 for i in range(n_occupied)]\n",
    "        beta_indices += [i * 2 + 1 + nelectrons for i in range(n_virtual)]\n",
    "        \n",
    "        for i in range(len(alpha_indices)):\n",
    "            p = alpha_indices[i]\n",
    "            for j in range(len(alpha_indices)):\n",
    "                q = alpha_indices[j]\n",
    "                    \n",
    "                qubit_op_dm = one_particle_op(p, q)\n",
    "                result = uccsd_circuit_vqe(spin_mult = 0, only_singles = False, only_doubles = False, qubits_num = qubits_num, \n",
    "                                   electron_count = nelectrons, optimize = False, theta = theta, \n",
    "                                   hamiltonian = qubit_op_dm, verbose = False)\n",
    "                dm_a[i, j] = result[0]\n",
    "                \n",
    "        for i in range(len(beta_indices)):\n",
    "            p = beta_indices[i]\n",
    "            for j in range(len(beta_indices)):\n",
    "                q = beta_indices[j]\n",
    "                    \n",
    "                qubit_op_dm = one_particle_op(p, q)\n",
    "                result = uccsd_circuit_vqe(spin_mult = 0, only_singles = False, only_doubles = False, qubits_num = qubits_num, \n",
    "                                   electron_count = nelectrons, optimize = False, theta = theta, \n",
    "                                   hamiltonian = qubit_op_dm, verbose = False)\n",
    "                dm_b[i, j] = result[0]\n",
    "\n",
    "        dm_ao = dm_a + dm_b\n",
    "        dm_ao = reduce(np.dot, (mf_pe.mo_coeff, dm_ao, mf_pe.mo_coeff.T))\n",
    "        \n",
    "        spin_pe_ham, E_pe, V_pe = pe_operator(dm_ao, mol, mf_pe.mo_coeff, water_pot, tolerance = 1e-12)\n",
    "\n",
    "        result = uccsd_circuit_vqe(spin_mult = 0, only_singles = False, only_doubles = False, qubits_num = qubits_num, \n",
    "                                   electron_count = nelectrons, optimize = False, theta = theta, \n",
    "                                   hamiltonian = spin_pe_ham, verbose = False)\n",
    "        edup = result[0]\n",
    "\n",
    "        # Check energy are the same\n",
    "        ovlp = mol.intor_symmetric('int1e_ovlp')\n",
    "        dm = np.einsum('pi,pq,qj->ij', ovlp, dm_ao, ovlp)\n",
    "        dm_mo = reduce(np.dot, (mf_pe.mo_coeff.T, dm, mf_pe.mo_coeff))\n",
    "        edup_np = np.einsum('ij,ji->', V_pe, dm_mo)\n",
    "        print('Energy from observe call= {:15g} & Energy from np einsum= {:15g}' .format(edup, edup_np), flush=True)\n",
    "\n",
    "        e_current = vqe_result[0] - edup + E_pe\n",
    "        dE = np.abs(e_current - e_last)\n",
    "        e_last = e_current\n",
    "\n",
    "        print('Cycle {:d}  E_tot = {:.15g}  dE = {:g}' .format(cycle, e_last, dE), flush=True)\n",
    "        print('Polarizable embedding energy: ', E_pe, flush=True)\n",
    "        print('Expectation value of polarizable operator: ', edup, flush=True)\n",
    "        print('\\n')\n",
    "\n",
    "        cycle+=1\n",
    "    else:\n",
    "        #raise ValueError(\"Unsuccessful optimization.\")\n",
    "        print(\"Unsuccessful optimization. Restart the optimization.\", flush=True)\n",
    "        \n",
    "        pass\n",
    "                \n",
    "        \n",
    "# Compute the RDM in molecular orbitals and natural orbital occupation number.\n",
    "ovlp = mol.intor_symmetric('int1e_ovlp')\n",
    "dm = np.einsum('pi,pq,qj->ij', ovlp, dm_ao, ovlp)\n",
    "dm_mo = reduce(np.dot, (mf_pe.mo_coeff.T, dm, mf_pe.mo_coeff))\n",
    "noon, U = np.linalg.eigh(dm_mo)\n",
    "noon = np.flip(noon)\n",
    "  \n",
    "print('\\n')\n",
    "print('Final result: ')\n",
    "print('Polarizable embedding energy: ', E_pe)\n",
    "print('Expectation value of polarizable operator: ', edup, flush=True)\n",
    "print('PE-VQE-UCCSD energy (L-BFGS-B)= ', e_last)\n",
    "print('Natural orbitals occupation number: ', noon)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01ac7926",
   "metadata": {},
   "source": [
    "### Example 2: NH3 with 46 water molecule using active space. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "88d53596",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "overwrite output file: qchem/NH3-pyscf.log\n",
      "[Pyscf] Total number of electrons =  10\n",
      "[Pyscf] Total number of orbitals =  15\n",
      "[Pyscf] Total HF energy with solvent: -56.460544146152934\n",
      "[Pyscf] Polarizable embedding energy from HF:  -0.2987297397877397\n",
      "[pyscf] CCSD energy of the active space (using molecular orbitals) with solvent=  -56.47023410365107\n",
      "[Pyscf] Polarizable embedding energy from CCSD:  -0.2986368707105786\n",
      "        Running solver for induced moments.\n",
      "0        --- Norm: 0.572427861637\n",
      "1        --- Norm: 0.194328219249\n",
      "        --- Turning on DIIS. ---\n",
      "2        --- Norm: 0.063478479823\n",
      "3        --- Norm: 0.007602903988\n",
      "4        --- Norm: 0.002247104251\n",
      "5        --- Norm: 0.000460072303\n",
      "6        --- Norm: 0.000114086983\n",
      "7        --- Norm: 0.000023144988\n",
      "8        --- Norm: 0.000005162818\n",
      "9        --- Norm: 0.000001149278\n",
      "10        --- Norm: 0.000000248690\n",
      "11        --- Norm: 0.000000069750\n",
      "12        --- Norm: 0.000000018086\n",
      "13        --- Norm: 0.000000001401\n"
     ]
    }
   ],
   "source": [
    "from qchem.classical_pyscf import get_mol_pe_hamiltonian\n",
    "from qchem.hamiltonian import jordan_wigner_fermion, jordan_wigner_pe\n",
    "\n",
    "\n",
    "geometry_NH3 = 'qchem/NH3.xyz'\n",
    "water_pot = 'qchem/46_water.pot'\n",
    "\n",
    "\n",
    "molecular_data=get_mol_pe_hamiltonian(xyz=geometry_NH3, potfile=water_pot, spin=0, charge=0, basis='631g', \\\n",
    "                                  nele_cas=6, norb_cas=6, ccsd=True, verbose=True)\n",
    "\n",
    "obi_mol = molecular_data[0]\n",
    "tbi_mol = molecular_data[1]\n",
    "e_core = molecular_data[2]\n",
    "obi_pe  =molecular_data[3]\n",
    "nelectrons = molecular_data[4]\n",
    "norbitals = molecular_data[5]\n",
    "\n",
    "qubits_num = 2 * norbitals\n",
    "\n",
    "\n",
    "spin_mol_ham = jordan_wigner_fermion(obi_mol, tbi_mol, e_core)\n",
    "spin_pe_ham = jordan_wigner_pe(obi_pe)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ff9b3f20",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "converged SCF energy = -56.460544146164\n"
     ]
    }
   ],
   "source": [
    "from pyscf import gto, scf, mcscf, solvent\n",
    "\n",
    "mol = gto.M(\n",
    "    atom = geometry_NH3,\n",
    "    spin = 0,\n",
    "    charge = 0,\n",
    "    basis = '631g'\n",
    ")\n",
    "mf_pe = scf.RHF(mol)\n",
    "mf_pe = solvent.PE(mf_pe, water_pot).run()\n",
    "mc = mcscf.CASCI(mf_pe, norbitals, nelectrons)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e809d4d",
   "metadata": {},
   "source": [
    "Note: for more accurate result, user needs to change `conv_tol` to much smaller value, e.g., 1e-7. In addition, the tolerance for the classical optimizer `vqe_tol` should be modified to smaller value, e.g., 1e-7."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c10e0dbc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of parameters: 18 singles, 99 doubles, 117 total\n",
      "Optimizer exited successfully:  True\n",
      "VQE Energy: -56.07492192212563\n",
      "VQE converged: True\n",
      "        Running solver for induced moments.\n",
      "0        --- Norm: 0.560904998895\n",
      "1        --- Norm: 0.195464440561\n",
      "        --- Turning on DIIS. ---\n",
      "2        --- Norm: 0.080909047471\n",
      "3        --- Norm: 0.014193006098\n",
      "4        --- Norm: 0.004526458635\n",
      "5        --- Norm: 0.001256746907\n",
      "6        --- Norm: 0.000476278998\n",
      "7        --- Norm: 0.000120262763\n",
      "8        --- Norm: 0.000033369845\n",
      "9        --- Norm: 0.000007548513\n",
      "10        --- Norm: 0.000001932468\n",
      "11        --- Norm: 0.000000513677\n",
      "12        --- Norm: 0.000000251862\n",
      "13        --- Norm: 0.000000062997\n",
      "14        --- Norm: 0.000000031247\n",
      "15        --- Norm: 0.000000008129\n",
      "Cycle 1  E_tot = -56.4704454234038  dE = 56.4704\n",
      "Polarizable embedding energy:  -0.29871334470788147\n",
      "Expectation value of polarizable operator:  0.09681015657024516\n",
      "\n",
      "\n",
      "Optimizer exited successfully:  True\n",
      "VQE Energy: -56.074941291297016\n",
      "VQE converged: True\n",
      "        Running solver for induced moments.\n",
      "0        --- Norm: 0.578438887906\n",
      "1        --- Norm: 0.186313099889\n",
      "        --- Turning on DIIS. ---\n",
      "2        --- Norm: 0.068418421402\n",
      "3        --- Norm: 0.012730202853\n",
      "4        --- Norm: 0.004094551612\n",
      "5        --- Norm: 0.000999816754\n",
      "6        --- Norm: 0.000319250864\n",
      "7        --- Norm: 0.000077598802\n",
      "8        --- Norm: 0.000023880904\n",
      "9        --- Norm: 0.000008037391\n",
      "10        --- Norm: 0.000001859332\n",
      "11        --- Norm: 0.000000532803\n",
      "12        --- Norm: 0.000000133206\n",
      "13        --- Norm: 0.000000032446\n",
      "14        --- Norm: 0.000000009283\n",
      "Cycle 2  E_tot = -56.4704648912704  dE = 1.94679e-05\n",
      "Polarizable embedding energy:  -0.2987133198747535\n",
      "Expectation value of polarizable operator:  0.09681028009861878\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "Final result: \n",
      "Polarizable embedding energy:  -0.2987133198747535\n",
      "Expectation value of polarizable operator:  0.09681028009861878\n",
      "PE-VQE-UCCSD energy (L-BFGS-B)=  -56.47046489127038\n",
      "Natural orbitals occupation number:  [ 2.00000000e+00  2.00000000e+00  1.99935887e+00  1.99541659e+00\n",
      "  1.99508319e+00  4.74450015e-03  4.36742635e-03  1.08093504e-03\n",
      "  4.89965694e-15  1.85085120e-15  1.30851294e-15  5.07383839e-16\n",
      " -1.05732651e-15 -1.59225347e-15 -6.12772938e-15]\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "from functools import reduce\n",
    "from qchem.uccsd import uccsd_parameter_size\n",
    "from qchem.uccsd_vqe import uccsd_circuit_vqe\n",
    "from qchem.particle_operator import one_particle_op\n",
    "from qchem.PEoperator import pe_operator_as\n",
    "from qchem.uccsd_init_param import get_parameters\n",
    "\n",
    "# Get the number of parameters\n",
    "singles, doubles, total = uccsd_parameter_size(nelectrons, qubits_num, spin_mult = 0)\n",
    "print(f\"Number of parameters: {singles} singles, {doubles} doubles, {total} total\")\n",
    "\n",
    "#theta = get_parameters(xyz = geometry_NH3, spin = 0,charge = 0, basis = '631g', nele_cas = 6, norb_cas = 6, \\\n",
    "#                           ccsd = True, without_solvent = False, potfile = water_pot)\n",
    "theta = np.zeros(total, dtype=np.float64)\n",
    "\n",
    "e_last = 0.0\n",
    "conv_tol = 1e-2\n",
    "dE = 1.0\n",
    "cycle = 1\n",
    "vqe_tol = 1e-3\n",
    "\n",
    "while dE > conv_tol:\n",
    "#for i in range (1):\n",
    "    spin_ham = spin_mol_ham + spin_pe_ham\n",
    "    \n",
    "    vqe_result = uccsd_circuit_vqe(spin_mult = 0, only_singles = False, only_doubles = False, qubits_num = qubits_num, \n",
    "                                   electron_count = nelectrons, optimize = True, theta = theta, \n",
    "                                   hamiltonian = spin_ham, method = 'L-BFGS-B', vqe_tol = vqe_tol, verbose = False)\n",
    "    print(f\"VQE Energy: {vqe_result[0]}\")\n",
    "    #print(f\"VQE parameters: {vqe_result[1]}\")\n",
    "    print(f\"VQE converged: {vqe_result[2]}\")\n",
    "    \n",
    "    theta = vqe_result[1]\n",
    "    \n",
    "    if vqe_result[2]:\n",
    "\n",
    "        dm_a = np.zeros((qubits_num // 2, qubits_num // 2))\n",
    "        dm_b = np.zeros((qubits_num // 2, qubits_num // 2))\n",
    "        dm = np.zeros((qubits_num, qubits_num))\n",
    "\n",
    "        n_spatial_orbitals = qubits_num // 2\n",
    "    \n",
    "        n_occupied = nelectrons // 2\n",
    "        n_virtual = n_spatial_orbitals - n_occupied\n",
    "\n",
    "        alpha_indices = [i * 2 for i in range(n_occupied)]\n",
    "        alpha_indices += [i * 2 + nelectrons for i in range(n_virtual)]\n",
    "\n",
    "        beta_indices = [i * 2 + 1 for i in range(n_occupied)]\n",
    "        beta_indices += [i * 2 + 1 + nelectrons for i in range(n_virtual)]\n",
    "        \n",
    "        for i in range(len(alpha_indices)):\n",
    "            p = alpha_indices[i]\n",
    "            for j in range(len(alpha_indices)):\n",
    "                q = alpha_indices[j]\n",
    "                    \n",
    "                qubit_op_dm = one_particle_op(p, q)\n",
    "                result = uccsd_circuit_vqe(spin_mult = 0, only_singles = False, only_doubles = False, qubits_num = qubits_num, \n",
    "                                   electron_count = nelectrons, optimize = False, theta = theta, \n",
    "                                   hamiltonian = qubit_op_dm, verbose = False)\n",
    "                dm_a[i, j] = result[0]\n",
    "                \n",
    "        for i in range(len(beta_indices)):\n",
    "            p = beta_indices[i]\n",
    "            for j in range(len(beta_indices)):\n",
    "                q = beta_indices[j]\n",
    "                    \n",
    "                qubit_op_dm = one_particle_op(p, q)\n",
    "                result = uccsd_circuit_vqe(spin_mult = 0, only_singles = False, only_doubles = False, qubits_num = qubits_num, \n",
    "                                   electron_count = nelectrons, optimize = False, theta = theta, \n",
    "                                   hamiltonian = qubit_op_dm, verbose = False)\n",
    "                dm_b[i, j] = result[0]\n",
    "\n",
    "        mocore = mf_pe.mo_coeff[:, :mc.ncore]\n",
    "        dm_core = np.dot(mocore, mocore.conj().T) * 2\n",
    "        casdm = dm_a + dm_b\n",
    "        mocas = mf_pe.mo_coeff[:, mc.ncore:mc.ncore + mc.ncas]\n",
    "        \n",
    "        # RDM in atomic orbitals\n",
    "        dm = dm_core + reduce(np.dot, (mocas, casdm, mocas.conj().T))\n",
    "\n",
    "        \n",
    "        spin_pe_hamiltonian, E_pe, V_pe, V_pe_cas= pe_operator_as(dm,mol,mf_pe.mo_coeff,water_pot,mc)\n",
    "\n",
    "        result = uccsd_circuit_vqe(spin_mult = 0, only_singles = False, only_doubles = False, qubits_num = qubits_num, \n",
    "                                   electron_count = nelectrons, optimize = False, theta = theta, \n",
    "                                   hamiltonian = spin_pe_ham, verbose = False)\n",
    "        edup = result[0]\n",
    "\n",
    "        e_current = vqe_result[0] - edup + E_pe\n",
    "        dE = np.abs(e_current - e_last)\n",
    "        e_last = e_current\n",
    "\n",
    "        print('Cycle {:d}  E_tot = {:.15g}  dE = {:g}' .format(cycle, e_last, dE), flush=True)\n",
    "        print('Polarizable embedding energy: ', E_pe, flush=True)\n",
    "        print('Expectation value of polarizable operator: ', edup, flush=True)\n",
    "        print('\\n')\n",
    "\n",
    "        cycle+=1\n",
    "    else:\n",
    "        print(\"Unsuccessful optimization. Restart the optimization.\", flush=True)\n",
    "        \n",
    "        pass\n",
    "            \n",
    "\n",
    "# convert dm from ao to mo (D_MO= C.T S D_AO S C ) and compute natural orbital occupation number.\n",
    "ovlp = mol.intor_symmetric('int1e_ovlp')\n",
    "dm = np.einsum('pi,pq,qj->ij', ovlp, dm, ovlp)\n",
    "dm_mo = reduce(np.dot, (mf_pe.mo_coeff.T, dm, mf_pe.mo_coeff))\n",
    "noon, U = np.linalg.eigh(dm_mo)\n",
    "noon = np.flip(noon)\n",
    "  \n",
    "print('\\n')\n",
    "print('Final result: ')\n",
    "print('Polarizable embedding energy: ', E_pe)\n",
    "print('Expectation value of polarizable operator: ', edup, flush=True)\n",
    "print('PE-VQE-UCCSD energy (L-BFGS-B)= ', e_last)\n",
    "print('Natural orbitals occupation number: ', noon)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "0fca85a5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CUDA-Q Version cu12-latest (https://github.com/NVIDIA/cuda-quantum 615d91910310605c83ea59f6afe6e7ae6dfecd28)\n"
     ]
    }
   ],
   "source": [
    "print(cudaq.__version__)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
